PR #509 Review: fix: async engine side-effect column propagation and collision resolution

Summary

This PR fixes two bugs in the async engine's handling of @custom_column_generator side-effect columns that prevented multi-stage pipelines (e.g., Anonymizer's detection workflow) from running with DATA_DESIGNER_ASYNC_ENGINE=1:

1. ExecutionGraph.set_side_effect() — changed from last-writer-wins to first-writer-wins semantics, preventing false DAGCircularDependencyError when multiple generators declare the same side-effect column at different pipeline stages.
2. AsyncTaskScheduler.__init__ — now includes side-effect columns in _instance_to_columns so their values are written to RowGroupBufferManager and available to downstream prompt templates.

Files changed: 3 (37 additions, 3 deletions)

Findings

Correctness

execution_graph.py — first-writer-wins (lines 107-116): The change is clean and correct. The set_side_effect method now guards with if side_effect_col not in self._side_effect_map, which implements first-writer-wins. Since ExecutionGraph.create() iterates column_configs in declared order (line 56-67), and the docstring says this matches sync engine semantics, this is the right fix. The previous last-writer-wins behavior would silently re-map a side-effect to a later producer, which could create a dependency edge back to an earlier stage and trigger a false cycle detection.

async_scheduler.py — side-effect propagation (lines 130-139): The fix correctly adds side-effect columns to _instance_to_columns. The seen_cols set prevents the same side-effect column from being added to multiple generators' output lists when there's a collision (consistent with the first-writer-wins semantics in the graph). The getattr(gen.config, "side_effect_columns", []) pattern is safe — configs that don't define side_effect_columns will return [].

Observations

1. Two-loop pattern in async_scheduler.py (lines 132-139): The code uses two separate loops over generators.items() — the first populates primary columns, the second adds side-effects. This is intentional and correct: the seen_cols set must be fully populated with all primary column names before side-effect deduplication runs, otherwise a side-effect column that shadows a primary column could slip through. Good design choice.

2. _instance_to_columns usage is comprehensive. I checked all 6 usage sites of _instance_to_columns in async_scheduler.py (lines ~358, ~379, ~621, ~666, ~768, ~794). All sites iterate output_cols and check if col in result_df.columns or if col in result before writing. This means adding side-effect columns to the list is safe — if a generator doesn't produce a side-effect value for a given invocation, it's simply skipped.

3. No logging for silently ignored collisions. When set_side_effect silently drops a second registration, there's no debug-level log. This is fine for correctness, but a logger.debug(...) could help users troubleshoot pipelines. Minor — not a blocker.

Testing

The new test test_side_effect_collision_first_writer_wins is well-structured: it registers two producers for the same side-effect, verifies the first wins, and checks downstream edge resolution. The existing test_side_effect_name_collision_prefers_real_column test (already passing) provides complementary coverage for the real-column-wins-over-side-effect case.

Gap: There is no unit test for the AsyncTaskScheduler side-effect propagation change. Testing this would require mocking generators with config.side_effect_columns, which is more of an integration concern. The PR description mentions 1747 engine tests pass, which likely covers this path indirectly. Low risk, but worth noting.

Style / Conventions

• Code follows the project's absolute import convention.
• Type annotations are present on all new variables (dict[int, list[str]], set[str]).
• Docstring on set_side_effect clearly explains the first-writer-wins contract and its rationale.
• No unnecessary changes or scope creep.

Verdict

LGTM. This is a focused, well-reasoned bug fix. Both changes are correct and consistent with each other. The test coverage for ExecutionGraph is solid. The only minor gap is the lack of a direct unit test for the AsyncTaskScheduler propagation change, but the overall test suite passing mitigates this risk. Ship it.

Greptile Summary

This PR fixes two bugs in the async engine's handling of @custom_column_generator side-effect columns. First, AsyncTaskScheduler.__init__ now maintains two separate dicts — _gen_instance_to_columns (graph-registered columns only, used for completion tracking) and _gen_instance_to_columns_including_side_effects (extends the former with side-effect columns, used for buffer writes). This directly fixes the KeyError in CompletionTracker described in the prior thread. Second, both ExecutionGraph.set_side_effect() and topologically_sort_column_configs() now raise ConfigCompilationError instead of silently overwriting when two different producers try to claim the same side-effect column name, converting a confusing false DAGCircularDependencyError into a clear diagnostic. Tests cover all three changed code paths with good isolation.

Important Files Changed

sequenceDiagram
    participant S as AsyncTaskScheduler
    participant R as _run_from_scratch / _run_cell / _run_batch
    participant BM as RowGroupBufferManager
    participant CT as CompletionTracker

    Note over S: __init__: build two dicts
    S->>S: _gen_instance_to_columns (real columns only)
    S->>S: _gen_instance_to_columns_including_side_effects (real + side-effect columns)

    S->>R: execute task (task.column)
    R->>R: generator.agenerate(...)
    R->>BM: update_batch / update_cell using _gen_instance_to_columns_including_side_effects

    S->>CT: mark_row_range_complete / mark_cell_complete using _gen_instance_to_columns (NO side-effect columns)

_{Reviews (7): Last reviewed commit: "fix: reject duplicate side-effect produc..." | Re-trigger Greptile}

PR #509 Review: fix: async engine side-effect column propagation and collision resolution

Summary

This PR fixes two bugs in the async engine's handling of @custom_column_generator side-effect columns that prevented multi-stage pipelines (e.g., Anonymizer's detection workflow) from running with DATA_DESIGNER_ASYNC_ENGINE=1 (fixes #508):

1. ExecutionGraph.set_side_effect() — changed from last-writer-wins to first-writer-wins semantics, with a warning log on collision. Prevents false DAGCircularDependencyError when multiple generators declare the same side-effect column at different pipeline stages.
2. AsyncTaskScheduler — now maintains a separate _gen_instance_to_columns_including_side_effects mapping so side-effect column values are written to RowGroupBufferManager and available to downstream prompt templates, while keeping completion tracking (_gen_instance_to_columns) free of unregistered columns.

Files changed: 4 (+117/−19)

Findings

Correctness

execution_graph.py — first-writer-wins (set_side_effect, lines 110-126): The change is clean and correct. The method now guards with if side_effect_col not in self._side_effect_map, implementing first-writer-wins. Since ExecutionGraph.create() iterates column_configs in declared order, and the docstring states this matches sync engine semantics (where earlier consumers see the first producer's value on the shared mutable row dict), this is the right fix. The previous last-writer-wins behavior silently re-mapped a side-effect to a later producer, creating incorrect dependency edges and false cycle detection. The added logger.warning on collision is a good addition for debuggability.

async_scheduler.py — dual-map approach (lines 134-151): The separation of _gen_instance_to_columns (completion tracking) from _gen_instance_to_columns_including_side_effects (buffer writes) is sound. Side-effect columns are not registered in the ExecutionGraph and would cause KeyError in CompletionTracker if included in the completion map. The seen_cols set prevents the same side-effect column from being added to multiple generators' output lists (consistent with first-writer-wins in the graph). The getattr(gen.config, "side_effect_columns", []) pattern safely handles configs without that attribute.

Two-loop pattern (lines 139-150): The code correctly uses two separate loops over generators.items() — first populates primary columns, then adds side-effects. This ensures seen_cols is fully populated with all primary column names before side-effect deduplication, preventing a side-effect from shadowing a real column from a different generator. Intentional and correct.

Usage-site audit — _gen_instance_to_columns vs _gen_instance_to_columns_including_side_effects:

• _gen_instance_to_columns is used in _salvage_rounds (lines 369, 390), _dispatch_seeds (line 633), and _execute_task_inner_impl (line 678) — all for dispatch/completion logic. Correct: side-effect columns should not participate in task scheduling.
• _gen_instance_to_columns_including_side_effects is used in _run_from_scratch (line 783), _run_cell (line 809), and _run_batch (line 833) — all for buffer writes. Correct: these are the points where result data is written back, and side-effect values must be included. All sites guard with if col in result_df.columns or if col in result, so missing side-effect values are safely skipped.

Minor Observations

1. [Low] Verbose set comprehension (async_scheduler.py:144): {col for col in generators} is equivalent to set(generators). Purely stylistic, not a bug.

2. [Low] Long line (async_scheduler.py:145): The dict comprehension line is ~107 characters. If this exceeds the project's ruff line-length config, it would need a break. If CI passes, it's fine.

3. [Info] Silent skip on side-effect/real-column overlap (async_scheduler.py:148): If a side-effect column name matches a real column from a different generator, the if side_effect_col not in seen_cols check silently skips it. This is correct (real columns take precedence, matching resolve_side_effect logic), but there's no warning logged for this case, unlike the graph collision. A debug log here could aid troubleshooting complex pipelines. Not a blocker.

Testing

test_side_effect_collision_first_writer_wins (test_execution_graph.py) — Well-structured: registers two producers for the same side-effect, verifies the first wins, checks the warning log output, and validates downstream dependency resolution via resolve_side_effect + add_edge. Good complementary coverage alongside existing test_side_effect_name_collision_prefers_real_column.

test_side_effect_columns_separated_from_completion_tracking (test_async_scheduler.py) — Directly tests the dual-map invariant: side-effect columns present in _gen_instance_to_columns_including_side_effects but absent from _gen_instance_to_columns. Uses object.__setattr__ to inject a mock config with side_effect_columns, which is a pragmatic approach for testing init logic. Good coverage of the key invariant.

Gap: No E2E test exercising the full async pipeline with side-effect columns end-to-end (acknowledged in PR description). Given the two unit tests cover the specific root causes and the full test suite passes (1747 engine tests), this is acceptable for merge. A follow-up E2E test would strengthen confidence for the Anonymizer use case.

Style / Conventions

• Absolute imports throughout, consistent with project rules.
• Type annotations on all new variables (dict[int, list[str]], set[str]).
• from __future__ import annotations present in modified files.
• Docstring on set_side_effect clearly explains the first-writer-wins contract and rationale.
• Inline comments in __init__ explain the purpose of each map. Verbose attribute name (_gen_instance_to_columns_including_side_effects) is long but self-documenting.
• No scope creep — changes are tightly focused on the two bugs.

Verdict

LGTM. This is a focused, well-reasoned bug fix. Both changes are correct and mutually consistent. The dual-map design in the scheduler cleanly separates completion tracking from buffer writes without disrupting existing async semantics. The first-writer-wins policy in the graph faithfully reproduces sync engine behavior. Test coverage is solid for both fixes. Minor style nits (verbose set comprehension, long line) are non-blocking. Ship it.